import os
import numpy as np
import torch
from torch.utils.data import Dataset
from torch.utils.data.sampler import (
    Sampler, SequentialSampler, RandomSampler, SubsetRandomSampler,
    WeightedRandomSampler)
from PIL import Image
from collections import defaultdict


label_dict = {
    'A': 0,
    'B': 1,
    'C': 2,
    'D': 3,
    'E': 4,
    'F': 5,
    'G': 6,
    'H': 7,
    'I': 8,
    'J': 9,
    'K': 10,
}


class Plane11Data(Dataset):
    def __init__(self, datadir, transform_type=None):
        self.labelname = label_dict
        self.abs_datadir = os.path.abspath(datadir)
        self.datainfo = self.getinfo()
        self.transform = transform_type

    def __getitem__(self, index):
        path_img, label = self.datainfo[index]
        img = Image.open(path_img).convert('RGB')
        if self.transform is not None:
            img = self.transform(img)
        return img, label

    def __len__(self):
        return len(self.datainfo)

    def getinfo(self):
        datalst = []
        for root, dirs, files in os.walk(self.abs_datadir):
            for filename in files:
                if filename.endswith('.png'):
                    cls_name = filename.split('_')[0]
                    path_png = os.path.abspath(os.path.join(root, filename))
                    datalst.append((path_png, self.labelname[cls_name]))
        # print(datalst)
        return datalst


class RandomIdentitySampler(Sampler):
    def __init__(self, data_source, num_instances=2):
        self.data_source = data_source
        self.num_instances = num_instances
        print('num_instances:', self.num_instances)
        
        self.index_dic = defaultdict(list)  # 按“主题”整理，存放该“主题”的所有“页码”
        
        for index, (_, pid) in enumerate(data_source):
            self.index_dic[pid].append(index)
        self.pids = list(self.index_dic.keys())
        self.num_pids = len(self.pids)

    def __len__(self):
        return self.num_pids * self.num_instances

    def __iter__(self):
        indices = torch.randperm(self.num_pids)  # 打乱pid顺序
        ret = []
        for i in indices:
            pid = self.pids[i]  # 按照乱序抽签号码来随机抽取一个pid
            t = self.index_dic[pid]  # t存放该pid的所有页码
            if len(t) >= self.num_instances:  # 页码充足，则避免采样重复
                t = np.random.choice(t, size=self.num_instances, replace=False)
            else:  # 页码不够，则允许重复采样
                t = np.random.choice(t, size=self.num_instances, replace=True)
            ret.extend(t)
        return iter(ret)
        # 竟然没有返回相应pid的信息
